from glob import glob
import numpy as np
from tqdm import tqdm
import pandas as pd

total_list = []

path = glob("E:/glm4_xi_out/*")
for one in tqdm(path):
    with open(one, "r", encoding="utf-8") as f:
        data = f.read()
    if len(data):
        data_list = [i.split("\n") for i in tqdm(data.split("\n\n")) if "以下" not in i]
        one_data_list = []

        for i in tqdm(np.hstack(data_list)):

            try:
                if i:
                    if i == "---":
                        continue
                    if i == "```":
                        continue

                    if i[-1] == "：":
                        continue

                    if "（" in i.split("：")[0] and "角色" in i.split("：")[0]:
                        new_data = i[i.index("（") + 1:]
                        new_data = new_data[:new_data.index("）")] + new_data[new_data.index("）") + 1:]
                        one_data_list.append(new_data)
                    else:
                        one_data_list.append(i)
            except:
                continue
        total_list.append(one_data_list)
# 拆分 512
new_total_list = []
token_list = ["<|p_{}|>".format(i) for i in range(515)]
voc_dict = dict()
for one in tqdm(total_list):
    two_str = ""
    for two in one:
        if len(two_str + two) < 512:
            two_str += two
        else:
            new_two = ["<|aos|>"] + list(two_str) + ["<|bos|>"] + token_list[
                                                                  len(["<|aos|>"] + list(two_str) + ["<|bos|>"]):]
            new_total_list.append(
                new_two)
            for i in new_two:
                voc_dict[i] = ""
            two_str = ""

voc = ["<|pad|>"] + list(voc_dict.keys())
pd.to_pickle(voc, "voc_data.pandas_pickle")
voc_dict = {v: k for k, v in enumerate(voc)}
new_total_list = [[voc_dict[j] for j in i] for i in new_total_list]
pd.to_pickle(new_total_list, "train_data.pandas_pickle")
